from solvers.solverbase import SolverBase
from solvers.codina.codinabase import CodinaBase
from solvers.stabilizations import supg
from dolfin import *
import numpy

class CrankNicolson(SolverBase,CodinaBase):
    'Crank-Nicolson method with nonlinear-crosswind stabilization.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        CodinaBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        V = FunctionSpace(mesh,'CG',1)
        VV = VectorFunctionSpace(mesh,'CG',1)
        DG = FunctionSpace(mesh,'DG',0)

        def get_sigma(v,u_,u__,eps,dt,sigma,cache): # u_ is the solution @ current iteration, u__ is a solution of previous time level
            phiS = TestFunction(DG)  # universal test function   
            trial = TrialFunction(DG) # trial function for seminorm and norm of solution and residuum 
        
            solSemF = Function(DG) #this is where the values will go
            resF = Function(DG)
            solNormF = Function(DG)            
      
            if 'a' not in cache.keys():
                a = trial*phiS*dx                                             # a
                cache['a'] = assemble(a)
        #----
            solSeml = sqrt(inner(nabla_grad(u_),nabla_grad(u_)))*phiS*dx   # l1 
            b1 = assemble(solSeml)
            solve(cache['a'],solSemF.vector(),b1)
        #-----     
            resl = (u_-u__+dt*inner(v,nabla_grad(u_)))*phiS*dx                 # l2               #TODO consider adding time derivate
            b2 = assemble(resl)
            solve(cache['a'],resF.vector(),b2)
        #-----
            solNorml = sqrt(resF*resF)*phiS*dx                                   # l3
            b3 = assemble(solNorml)
            solve(cache['a'],solNormF.vector(),b3)
        #-----
 
            nc = mesh.num_cells()
            sigmaValues = numpy.zeros(nc)
            beta = 0.7
            for i in xrange(nc): # build the values
                c = Cell(mesh,i)
               
                diamK = c.diameter()  #elements diameter
                x = c.midpoint()
                
                norm = solNormF(x.x(),x.y())
                seminorm = solSemF(x.x(),x.y())
                
                if norm < 1E-16 or seminorm < 1E-16:
                    sigmaValues[i] = 0.
                else:
                    sigmaValues[i] = 0.5*max(0,beta - 2*eps*seminorm/norm/diamK)*diamK*norm/seminorm
                
            sigma.vector()[:] = sigmaValues
        #-----------------------------------------------------------------

        phi_ = interpolate(problem.phi_,V)
        phi = TrialFunction(V)
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)
        tau = Function(DG)
        tau.vector()[:] = supg.get_tau(v,eps=supg.diffusivity)

        sigma = Function(DG)                # stab. param for phi_n+1, changes during iteration
        sigma_ = Function(DG)               # stab. param for phi_, does not change during iteration

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0
        
        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                           ) # a way to realize swith on |v| == 0.
                        
        pMatrix = conditional(
                             gt(sqrt(inner(v,v)),1E-16),
                             Identity(2)-outer(v,v)/inner(v,v),
                             Constant(0.)*Identity(2),
                             ) # I - v X v
        
        a = phi*varphi*dx + 0.5*dt*inner(v,nabla_grad(phi))*varphi*dx +\
            tau*pert*phi*inner(v,nabla_grad(varphi))*dx +\
            0.5*dt*tau*pert*inner(dot(nabla_grad(phi),outer(v,v)),nabla_grad(varphi))*dx+\
            0.5*dt*sigma*inner(dot(nabla_grad(phi),pMatrix),nabla_grad(varphi))*dx

        l = phi_*varphi*dx - 0.5*dt*inner(v,nabla_grad(phi_))*varphi*dx+\
            tau*pert*phi_*inner(v,nabla_grad(varphi))*dx\
            - 0.5*dt*tau*pert*inner(dot(nabla_grad(phi_),outer(v,v)),nabla_grad(varphi))*dx\
            -0.5*dt*sigma*inner(dot(nabla_grad(phi_),pMatrix),nabla_grad(varphi))*dx
        # SOLO CODINA DOES NOT SEEM TO BE A GOOD IDEA FOR NONLINEAR METHOD 
        t = 0
        count = 0
        
        maxIter = 10
        tol = 1E-3
        
        cache = {}
        phi = Function(V)
        phi_m = Function(V)
        while t <= T:
            t += float(dt)
            count += 1
            iter = 0
            error = 1

            get_sigma(v,phi_,phi_,supg.diffusivity,float(dt),sigma,cache)    # get the sigma for first iteration with phi_0 = phi_(previous solution)
            while error > tol and iter < maxIter:
                iter += 1
                A = assemble(a)    
                L = assemble(l) 
                solve(A,phi.vector(),L)
                diff = phi.vector().array() - phi_m.vector().array()
                error = numpy.linalg.norm(diff,ord=numpy.Inf)
                phi_m.assign(phi)
                get_sigma(v,0.5*(phi_m+phi_),phi_,supg.diffusivity,float(dt),sigma,cache) 
                # get the sigma for the next iteration with phi_j = phi(current solution)

            phi_.assign(phi)
            self.update(problem,phi_,t,float(dt),count)

        problem.save_data(self.saveDir,str(self.CFL))
